/* stress.c: A program to put stress on a POSIX system (stress).
 *
 * Copyright (C) 2001,2002,2003,2004,2005,2006,2007,2008,2009,2010
 * Amos Waterland <apw@rossby.metr.ou.edu>
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc., 59
 * Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#include <ctype.h>
#include <errno.h>
#include <libgen.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <time.h>
#include <unistd.h>
#include <sys/wait.h>

#include "../config.h"
#include "stress.h"

/* By default, print all messages of severity info and above.  */
static int global_debug = 2;

/* Name of this program */
static char *global_progname = PACKAGE;

/* Implemention of runtime-selectable severity message printing.  */
#define dbg(OUT, STR, ARGS...) if (global_debug >= 3) \
	fprintf(stdout, "%s: dbug: [%lli] ", \
		global_progname, (long long)getpid()), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)
#define out(OUT, STR, ARGS...) if (global_debug >= 2) \
	fprintf(stdout, "%s: info: [%lli] ", \
		global_progname, (long long)getpid()), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)
#define wrn(OUT, STR, ARGS...) if (global_debug >= 1) \
	fprintf(stderr, "%s: WARN: [%lli] (%d) ", \
		global_progname, (long long)getpid(), __LINE__), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)
#define err(OUT, STR, ARGS...) if (global_debug >= 0) \
	fprintf(stderr, "%s: FAIL: [%lli] (%d) ", \
		global_progname, (long long)getpid(), __LINE__), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)

/* Implementation of check for option argument correctness.  */
#define assert_arg(A) \
          if (++i == argc || ((arg = argv[i])[0] == '-' && \
              !isdigit ((int)arg[1]) )) \
            { \
              err(stderr, "missing argument to option '%s'\n", A); \
              exit(1); \
            }

/* Prototypes for utility functions.  */
int usage(int status);
int version(int status);
long long atoll_s(const char *nptr);
long long atoll_b(const char *nptr);

/* Prototypes for worker functions.  */
int hogcpu(void);
int hogio(void);
int hogvm(long long bytes, long long stride, long long hang, int keep);
int hoghdd(long long bytes);

int
main(int argc, char **argv)
{
  int i, pid, children = 0, retval = 0;
  long starttime, stoptime, runtime, forks;

  /* Variables that indicate which options have been selected.  */
  int do_dryrun = 0;
  long long do_backoff = 3000;
  long long do_timeout = 0;
  long long do_cpu = 0;
  long long do_io = 0;
  long long do_vm = 0;
  long long do_vm_bytes = (256 * 1024 * 1024);
  /* (256 * 1024 * 1024) = 268435456 */
  long long do_vm_stride = 4096;
  long long do_vm_hang = -1;
  int do_vm_keep = 0;
  long long do_hdd = 0;
  long long do_hdd_bytes = (1024 * 1024 * 1024);
  /* (1024 * 1024 * 1024) = 1073741824 */

  /* Record our start time.  */
  if ((starttime = time (NULL)) == -1) {
      err(stderr, "failed to acquire current time: %s\n", strerror(errno));
      exit(1);
  }

  /* SuSv3 does not define any error conditions for this function.  */
  global_progname = basename(argv[0]);

  /* For portability, parse command line options without getopt_long(), and
   * instead just use a bunch of calls to strcmp():  */
  for ((i = 1); (i < argc); i++) {
      char *arg = argv[i];

      if ((strcmp(arg, "--help") == 0) || (strcmp(arg, "-?") == 0)) {
          usage(0);
	  } else if (strcmp(arg, "--version") == 0) {
          version(0);
	  } else if ((strcmp(arg, "--verbose") == 0) || (strcmp(arg, "-v") == 0)) {
          global_debug = 3;
	  } else if ((strcmp(arg, "--quiet") == 0) || (strcmp(arg, "-q") == 0)) {
          global_debug = 0;
	  } else if ((strcmp(arg, "--dry-run") == 0) || (strcmp(arg, "-n") == 0)) {
          do_dryrun = 1;
	  } else if (strcmp(arg, "--backoff") == 0) {
          assert_arg("--backoff");
          if (sscanf(arg, "%lli", &do_backoff) != 1) {
              err(stderr, "invalid number: %s\n", arg);
              exit(1);
		  }
          if (do_backoff < 0) {
              err(stderr, "invalid backoff factor: %lli\n", do_backoff);
              exit(1);
		  }
          dbg(stdout, "setting backoff coeffient to %llius\n", do_backoff);
	  } else if ((strcmp(arg, "--timeout") == 0) || (strcmp(arg, "-t") == 0)) {
          assert_arg("--timeout");
          do_timeout = atoll_s(arg);
          if (do_timeout <= 0) {
              err(stderr, "invalid timeout value: %llis\n", do_timeout);
              exit(1);
		  }
	  } else if ((strcmp(arg, "--cpu") == 0) || (strcmp(arg, "-c") == 0)) {
          assert_arg("--cpu");
          do_cpu = atoll_b(arg);
          if (do_cpu <= 0) {
              err(stderr, "invalid number of cpu hogs: %lli\n", do_cpu);
              exit(1);
		  }
	  } else if ((strcmp(arg, "--io") == 0) || (strcmp(arg, "-i") == 0)) {
          assert_arg("--io");
          do_io = atoll_b(arg);
          if (do_io <= 0) {
              err(stderr, "invalid number of io hogs: %lli\n", do_io);
              exit(1);
		  }
	  } else if ((strcmp(arg, "--vm") == 0) || (strcmp(arg, "-m") == 0)) {
          assert_arg("--vm");
          do_vm = atoll_b(arg);
          if (do_vm <= 0) {
              err(stderr, "invalid number of vm hogs: %lli\n", do_vm);
              exit(1);
		  }
	  } else if (strcmp(arg, "--vm-bytes") == 0) {
          assert_arg("--vm-bytes");
          do_vm_bytes = atoll_b(arg);
          if (do_vm_bytes <= 0) {
              err(stderr, "invalid vm byte value: %lli\n", do_vm_bytes);
              exit(1);
		  }
	  } else if (strcmp(arg, "--vm-stride") == 0) {
          assert_arg("--vm-stride");
          do_vm_stride = atoll_b(arg);
          if (do_vm_stride <= 0) {
              err(stderr, "invalid stride value: %lli\n", do_vm_stride);
              exit(1);
		  }
	  } else if (strcmp(arg, "--vm-hang") == 0) {
          assert_arg("--vm-hang");
          do_vm_hang = atoll_b(arg);
          if (do_vm_hang < 0) {
              err(stderr, "invalid value: %lli\n", do_vm_hang);
              exit(1);
		  }
	  } else if (strcmp(arg, "--vm-keep") == 0) {
          do_vm_keep = 1;
	  } else if ((strcmp(arg, "--hdd") == 0) || (strcmp(arg, "-d") == 0)) {
          assert_arg("--hdd");
          do_hdd = atoll_b(arg);
          if (do_hdd <= 0) {
              err(stderr, "invalid number of hdd hogs: %lli\n", do_hdd);
              exit(1);
		  }
	  } else if (strcmp(arg, "--hdd-bytes") == 0) {
          assert_arg("--hdd-bytes");
          do_hdd_bytes = atoll_b(arg);
          if (do_hdd_bytes <= 0) {
              err(stderr, "invalid hdd byte value: %lli\n", do_hdd_bytes);
              exit(1);
		  }
	  } else {
          err(stderr, "unrecognized option: %s\n", arg);
          exit(1);
	  }
  }

  /* Print startup message if we have work to do, bail otherwise.  */
  if (do_cpu + do_io + do_vm + do_hdd) {
      out(stdout, "dispatching hogs: %lli cpu, %lli io, %lli vm, %lli hdd\n",
		  do_cpu, do_io, do_vm, do_hdd);
  } else {
	  usage(0);
  }

  /* Round robin dispatch our worker processes.  */
  while ((forks = (do_cpu + do_io + do_vm + do_hdd))) {
      long long backoff, timeout = 0;

      /* Calculate the backoff value so we get good fork throughput.  */
      backoff = (do_backoff * forks);
      dbg(stdout, "using backoff sleep of %llius\n", backoff);

      /* If we are supposed to respect a timeout, calculate it.  */
      if (do_timeout) {
          long long currenttime;

          /* Acquire current time.  */
          if ((currenttime = time(NULL)) == -1) {
              perror("error acquiring current time");
              exit(1);
		  }

          /* Calculate timeout based on current time.  */
          timeout = (do_timeout - (currenttime - starttime));

          if (timeout > 0) {
              dbg(stdout, "setting timeout to %llis\n", timeout);
		  } else {
              wrn(stderr, "used up time before all workers dispatched\n");
              break;
		  }
	  }

      if (do_cpu) {
          switch (pid = fork()) {
            case 0:            /* child */
              alarm(timeout);
              usleep(backoff);
			  if (do_dryrun) {
                exit(0);
			  }
              exit(hogcpu());
            case -1:           /* error */
              err(stderr, "fork failed: %s\n", strerror(errno));
              break;
            default:           /* parent */
              dbg(stdout, "--> hogcpu worker %lli [%i] forked\n",
				  do_cpu, pid);
              ++children;
		  }
          --do_cpu;
	  }

      if (do_io) {
          switch (pid = fork()) {
            case 0:            /* child */
              alarm(timeout);
              usleep(backoff);
			  if (do_dryrun) {
                exit(0);
			  }
              exit(hogio());
            case -1:           /* error */
              err(stderr, "fork failed: %s\n", strerror(errno));
              break;
            default:           /* parent */
              dbg(stdout, "--> hogio worker %lli [%i] forked\n", do_io, pid);
              ++children;
		  }
          --do_io;
	  }

      if (do_vm) {
          switch (pid = fork()) {
            case 0:            /* child */
              alarm(timeout);
              usleep(backoff);
			  if (do_dryrun) {
                exit(0);
			  }
              exit(hogvm(do_vm_bytes, do_vm_stride, do_vm_hang, do_vm_keep));
            case -1:           /* error */
              err(stderr, "fork failed: %s\n", strerror(errno));
              break;
            default:           /* parent */
              dbg(stdout, "--> hogvm worker %lli [%i] forked\n", do_vm, pid);
              ++children;
		  }
          --do_vm;
	  }

      if (do_hdd) {
          switch (pid = fork()) {
            case 0:            /* child */
              alarm(timeout);
              usleep(backoff);
			  if (do_dryrun) {
                exit(0);
			  }
              exit(hoghdd(do_hdd_bytes));
            case -1:           /* error */
              err(stderr, "fork failed: %s\n", strerror(errno));
              break;
            default:           /* parent */
              dbg(stdout, "--> hoghdd worker %lli [%i] forked\n",
				  do_hdd, pid);
              ++children;
            }
          --do_hdd;
        }
    }

  /* Wait for our children to exit.  */
  while (children) {
      int status, ret;

      if ((pid = wait(&status)) > 0) {
          --children;

          if (WIFEXITED(status)) {
              if ((ret = WEXITSTATUS(status)) == 0) {
                  dbg(stdout, "<-- worker %i returned normally\n", pid);
			  } else {
                  err(stderr, "<-- worker %i returned error %i\n", pid, ret);
                  ++retval;
                  wrn (stderr, "now reaping child worker processes\n");
                  if (signal(SIGUSR1, SIG_IGN) == SIG_ERR) {
                    err(stderr, "handler error: %s\n", strerror(errno));
				  }
                  if (kill((-1 * getpid()), SIGUSR1) == -1) {
                    err(stderr, "kill error: %s\n", strerror(errno));
				  }
			  }
		  } else if (WIFSIGNALED(status)) {
              if ((ret = WTERMSIG(status)) == SIGALRM) {
                  dbg(stdout, "<-- worker %i signalled normally\n", pid);
			  } else if ((ret = WTERMSIG (status)) == SIGUSR1) {
                  dbg(stdout, "<-- worker %i reaped\n", pid);
			  } else {
                  err(stderr, "<-- worker %i got signal %i\n", pid, ret);
                  ++retval;
                  wrn (stderr, "now reaping child worker processes\n");
                  if (signal (SIGUSR1, SIG_IGN) == SIG_ERR) {
                    err(stderr, "handler error: %s\n", strerror(errno));
				  }
                  if (kill((-1 * getpid()), SIGUSR1) == -1) {
                    err(stderr, "kill error: %s\n", strerror(errno));
				  }
			  }
		  } else {
              err(stderr, "<-- worker %i exited abnormally\n", pid);
              ++retval;
		  }
	  } else {
          err(stderr, "error waiting for worker: %s\n", strerror(errno));
          ++retval;
          break;
	  }
  }

  /* Record our stop time.  */
  if ((stoptime = time(NULL)) == -1) {
      err(stderr, "failed to acquire current time\n");
      exit(1);
  }

  /* Calculate our runtime.  */
  runtime = (stoptime - starttime);

  /* Print final status message.  */
  if (retval) {
      err(stderr, "failed run completed in %lis\n", runtime);
  } else {
      out (stdout, "successful run completed in %lis\n", runtime);
  }

  exit(retval);
}

/* */
int
hogcpu(void)
{
	while (1) {
		sqrt(rand());
	}

  return 0;
}

int
hogio(void)
{
	while (1) {
		sync();
	}

  return 0;
}

int
hogvm(long long bytes, long long stride, long long hang, int keep)
{
  long long i;
  char *ptr = 0;
  char c;
  int do_malloc = 1;

  while (1) {
      if (do_malloc) {
          dbg(stdout, "allocating %lli bytes ...\n", bytes);
          if (!(ptr = (char *)malloc(bytes * sizeof(char)))) {
              err(stderr, "hogvm malloc failed: %s\n", strerror(errno));
              return 1;
		  }
          if (keep) {
			  do_malloc = 0;
		  }
	  }

      dbg(stdout, "touching bytes in strides of %lli bytes ...\n", stride);
      for ((i = 0); (i < bytes); (i += stride)) {
		  ptr[i] = 'Z';           /* Ensure that COW happens.  */
	  }

      if (hang == 0) {
          dbg(stdout, "sleeping forever with allocated memory\n");
          while (1) {
			  sleep(1024);
		  }
	  } else if (hang > 0) {
          dbg(stdout, "sleeping for %llis with allocated memory\n", hang);
          sleep(hang);
	  }

      for ((i = 0); (i < bytes); (i += stride)) {
          c = ptr[i];
          if (c != 'Z') {
              err(stderr, "memory corruption at: %p\n", (ptr + i));
              return 1;
		  }
	  }

      if (do_malloc) {
          free(ptr);
          dbg(stdout, "freed %lli bytes\n", bytes);
	  }
  }

  return 0;
}

int
hoghdd(long long bytes)
{
  long long i, j;
  int fd;
  int chunk = ((1024 * 1024) - 1);        /* Minimize slow writing.  */
  char buff[chunk];

  /* Initialize buffer with some random ASCII data.  */
  dbg(stdout, "seeding %d byte buffer with random data\n", chunk);
  for ((i = 0); (i < (chunk - 1)); i++) {
      j = rand();
      j = ((j < 0) ? -j : j);
      j %= 95;
      j += 32;
      buff[i] = j;
  }
  buff[i] = '\n';

  while (1) {
      char name[] = "./stress.XXXXXX";

      if ((fd = mkstemp (name)) == -1) {
          err(stderr, "mkstemp failed: %s\n", strerror(errno));
          return 1;
	  }

      dbg(stdout, "opened %s for writing %lli bytes\n", name, bytes);

      dbg(stdout, "unlinking %s\n", name);
      if (unlink (name) == -1) {
          err(stderr, "unlink of %s failed: %s\n", name, strerror(errno));
          return 1;
	  }

      dbg(stdout, "fast writing to %s\n", name);
      for ((j = 0); ((bytes == 0) || ((j + chunk) < bytes)); (j += chunk)) {
          if (write(fd, buff, chunk) == -1) {
              err(stderr, "write failed: %s\n", strerror(errno));
              return 1;
            }
        }

      dbg(stdout, "slow writing to %s\n", name);
      for (; ((bytes == 0) || (j < (bytes - 1))); j++) {
          if (write(fd, &buff[(j % chunk)], 1) == -1) {
              err(stderr, "write failed: %s\n", strerror(errno));
              return 1;
		  }
	  }
      if (write(fd, "\n", 1) == -1) {
          err(stderr, "write failed: %s\n", strerror(errno));
          return 1;
	  }
      ++j;

      dbg(stdout, "closing %s after %lli bytes\n", name, j);
      close(fd);
    }

  return 0;
}

/* Convert a string representation of a number with an optional size suffix
 * to a long long.
 */
long long
atoll_b(const char *nptr)
{
  int pos;
  char suffix;
  long long factor = 0;
  long long value;

  if ((pos = (strlen(nptr) - 1)) < 0) {
      err(stderr, "invalid string\n");
      exit(1);
  }

  switch (suffix = nptr[pos]) {
    case 'b':
    case 'B':
      factor = 0;
      break;
    case 'k':
    case 'K':
      factor = 10;
      break;
    case 'm':
    case 'M':
      factor = 20;
      break;
    case 'g':
    case 'G':
      factor = 30;
      break;
    default:
      if ((suffix < '0') || (suffix > '9')) {
          err(stderr, "unrecognized suffix: %c\n", suffix);
          exit(1);
	  }
  }

  if (sscanf(nptr, "%lli", &value) != 1) {
      err(stderr, "invalid number: %s\n", nptr);
      exit(1);
  }

  value = (value << factor);

  return value;
}

/* Convert a string representation of a number with an optional time suffix
 * to a long long.
 */
long long
atoll_s(const char *nptr)
{
  int pos;
  char suffix;
  long long factor;
  long long value;

  factor = 1;

  if ((pos = (strlen(nptr) - 1)) < 0) {
      err(stderr, "invalid string\n");
      exit(1);
  }

  switch (suffix = nptr[pos]) {
    case 's':
    case 'S':
      factor = 1;
      break;
    case 'm':
    case 'M':
      factor = 60;
      break;
    case 'h':
    case 'H':
      factor = (60 * 60);
	  /* (60 * 60) = 3600 */
      break;
    case 'd':
    case 'D':
      factor = (60 * 60 * 24);
	  /* (60 * 60 * 24) = 86400 */
      break;
    case 'y':
    case 'Y':
      factor = (60 * 60 * 24 * 365);
	  /* (60 * 60 * 24 * 365) = 31536000 */
      break;
    default:
      if ((suffix < '0') || (suffix > '9')) {
          err(stderr, "unrecognized suffix: %c\n", suffix);
          exit(1);
	  }
  }

  if (sscanf(nptr, "%lli", &value) != 1) {
      err(stderr, "invalid number: %s\n", nptr);
      exit(1);
  }

  value = (value * factor);

  return value;
}

int
version(int status)
{
  char *mesg = "%s %s\n";

  fprintf(stdout, mesg, global_progname, VERSION);

	if (status <= 0) {
		exit(-1 * status);
	}

  return 0;
}

int
usage(int status)
{
  char *mesg =
    "`%s' imposes certain types of compute stress on your system\n\n"
    "Usage: %s [OPTION [ARG]] ...\n"
    " -?, --help         show this help statement\n"
    "     --version      show version statement\n"
    " -v, --verbose      be verbose\n"
    " -q, --quiet        be quiet\n"
    " -n, --dry-run      show what would have been done\n"
    " -t, --timeout N    timeout after N seconds\n"
    "     --backoff N    wait factor of N microseconds before work starts\n"
    " -c, --cpu N        spawn N workers spinning on sqrt()\n"
    " -i, --io N         spawn N workers spinning on sync()\n"
    " -m, --vm N         spawn N workers spinning on malloc()/free()\n"
    "     --vm-bytes B   malloc B bytes per vm worker (default is 256MB)\n"
    "     --vm-stride B  touch a byte every B bytes (default is 4096)\n"
    "     --vm-hang N    sleep N secs before free (default none, 0 is inf)\n"
    "     --vm-keep      redirty memory instead of freeing and reallocating\n"
    " -d, --hdd N        spawn N workers spinning on write()/unlink()\n"
    "     --hdd-bytes B  write B bytes per hdd worker (default is 1GB)\n\n"
    "Example: %s --cpu 8 --io 4 --vm 2 --vm-bytes 128M --timeout 10s\n\n"
    "Note: Numbers may be suffixed with s,m,h,d,y (time) or B,K,M,G (size).\n";

  fprintf(stdout, mesg, global_progname, global_progname, global_progname);

	if (status <= 0) {
		exit(-1 * status);
	}

  return 0;
}

/* EOF */
